import io
from collections import defaultdict

import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt


# painter.py
class TP:
    def __init__(self, t, pid, step, action):
        self.t = t
        self.pid = pid
        self.step = step
        self.action = action

    def to_dict(self):
        return {"t": self.t, "pid": self.pid, "step": self.step, "action": self.action}

    @classmethod
    def from_dict(cls, data):
        return cls(data["t"], data["pid"], data["step"], data["action"])


class ProcessTimelinePlot:
    def __init__(self, tp_vector):
        self.tp_vector = tp_vector
        self.process_data = defaultdict(list)
        self.current_step = 0
        self.processes = set()
        self.colors = {}
        self.time_points = set()
        self.segments = []
        self._process_data()

    def _process_data(self):
        segment_start = {}
        for tp in self.tp_vector:
            self.processes.add(tp.pid)
            self.process_data[tp.pid].append(tp)
            self.time_points.add(tp.t)

            if tp.action == "started":
                segment_start[(tp.pid, tp.step)] = tp.t
            elif tp.action == "finished":
                if (tp.pid, tp.step) in segment_start:
                    start_time = segment_start[(tp.pid, tp.step)]
                    end_time = tp.t
                    linestyle = "-" if tp.step in ("c1", "c2") else (0, (3, 1))
                    self.segments.append((start_time, end_time, tp.pid, linestyle))
                    del segment_start[(tp.pid, tp.step)]

        self.processes = sorted(self.processes)
        self.colors = {
            pid: plt.cm.tab20(i / len(self.processes))
            for i, pid in enumerate(self.processes)
        }
        self.time_points = sorted(self.time_points)

    def _plot_segment(self, segment, ax):
        start_time, end_time, pid, linestyle = segment
        pid_index = self.processes.index(pid)
        color = self.colors[pid]
        ax.plot(
            [start_time, end_time],
            [pid_index, pid_index],
            linestyle=linestyle,
            color=color,
            linewidth=2,
        )

    def step(self, path):
        if self.current_step < len(self.segments):
            fig, ax = plt.subplots()
            ax.set_yticks(range(len(self.processes)))
            ax.set_yticklabels(self.processes)
            ax.set_ylim(-0.5, len(self.processes) - 0.5)  # 调整y轴间距
            ax.set_xlabel("Time")
            ax.set_ylabel("Processes")
            ax.set_xticks(self.time_points)
            ax.grid(axis="x")

            for i in range(self.current_step + 1):
                self._plot_segment(self.segments[i], ax)

            self.current_step += 1
            plt.savefig(path)

            # buf = io.BytesIO()
            # plt.savefig(buf, format="png")
            # plt.close(fig)
            # buf.seek(0)
            # return buf.getvalue()

    def plot_full(self, path):
        fig, ax = plt.subplots()
        ax.set_yticks(range(len(self.processes)))
        ax.set_yticklabels(self.processes)
        ax.set_ylim(-0.5, len(self.processes) - 0.5)  # 调整y轴间距
        ax.set_xlabel("Time")
        ax.set_ylabel("Processes")
        ax.set_xticks(self.time_points)
        ax.grid(axis="x")

        for i, segment in enumerate(self.segments):
            self._plot_segment(segment, ax)
            plt.savefig(path + f"plot_step_{i}.png", format="png")

        plt.savefig(path + "plot_full.png", format="png")


# # Example usage:
# tp_vector = [
#     TP(0, "a", "c1", "started"),
#     TP(30, "a", "c1", "finished"),
#     TP(30, "a", "io", "started"),
#     TP(30, "b", "c1", "started"),
#     TP(70, "a", "io", "finished"),
#     TP(90, "b", "c1", "finished"),
#     TP(90, "b", "io", "started"),
#     TP(90, "a", "c2", "started"),
#     TP(100, "b", "io", "finished"),
#     TP(110, "a", "c2", "finished"),
#     TP(110, "b", "c2", "started"),
#     TP(140, "b", "c2", "finished"),
#     TP(140, "c", "c1", "started"),
#     TP(170, "c", "c1", "finished"),
#     TP(170, "c", "io", "started"),
#     TP(210, "c", "io", "finished"),
#     TP(210, "c", "c2", "started"),
#     TP(230, "c", "c2", "finished"),
# ]
# plotter = ProcessTimelinePlot(tp_vector)
# # plotter.step()  # To plot step by step
# # plotter.step()  # To plot step by step
# # plotter.step()  # To plot step by step
# plotter.plot_full("output/")  # To plot the full timeline
